import tkinter as tk
from tkinter import scrolledtext, ttk, messagebox
from scapy.all import sniff, Raw, hexdump
from scapy.layers.inet import IP, UDP, TCP, ICMP
import threading
from datetime import datetime
import psutil

class PacketSniffer:
    def __init__(self, root):
        self.root = root
        self.root.title("网络嗅探器V1.0")
        self.root.geometry("900x600")
        self.root.configure(bg="#F0F0F0")
        self.sniffing = False
        self.captured_packets = []
        self.selected_pid = None  # 存储选定的进程ID

        self.create_frames()
        self.nic_var = tk.StringVar(value=self.get_nic_list()[0])
        self.create_protocol_widgets()
        self.create_buttons()
        self.create_packet_display_frame()
        self.create_packet_details_frame()
        self.create_16_data_frame()

    def create_frames(self):
        self.protocol_frame = tk.Frame(self.root, bg="#F0F0F0")
        self.protocol_frame.pack(side=tk.TOP, fill=tk.X)

        self.main_frame = tk.Frame(self.root, bg="#F0F0F0")
        self.main_frame.pack(side=tk.TOP, fill=tk.BOTH, expand=True)

    def create_protocol_widgets(self):
        self.nic_label = tk.Label(self.protocol_frame, text="选择NIC:", bg="#F0F0F0")
        self.nic_label.pack(side=tk.LEFT, padx=(10, 2), pady=(15, 5))

        self.nic_options = self.get_nic_list()
        self.nic_menu = ttk.Combobox(self.protocol_frame, textvariable=self.nic_var, values=self.nic_options)
        self.nic_menu.pack(side=tk.LEFT, padx=5, pady=(15, 5))

        self.host_var = tk.StringVar()
        self.host_label = tk.Label(self.protocol_frame, text="监听IP:", bg="#F0F0F0")
        self.host_label.pack(side=tk.LEFT, padx=(10, 2), pady=(15, 5))

        self.host_entry = ttk.Entry(self.protocol_frame, textvariable=self.host_var, width=15)
        self.host_entry.pack(side=tk.LEFT, padx=5, pady=(15, 5))

        self.port_var = tk.StringVar()
        self.port_label = tk.Label(self.protocol_frame, text="Port:", bg="#F0F0F0")
        self.port_label.pack(side=tk.LEFT, padx=(10, 2), pady=(15, 5))

        self.port_entry = ttk.Entry(self.protocol_frame, textvariable=self.port_var, width=8)
        self.port_entry.pack(side=tk.LEFT, padx=5, pady=(15, 5))

        self.protocol_var = tk.StringVar(value="ALL")
        self.protocol_label = tk.Label(self.protocol_frame, text="协议:", bg="#F0F0F0")
        self.protocol_label.pack(side=tk.LEFT, padx=(10, 2), pady=(15, 5))

        self.protocol_options = ["ALL", "HTTP", "TCP", "UDP", "IPv4", "IPv6", "ICMP"]
        self.protocol_menu = ttk.Combobox(self.protocol_frame, textvariable=self.protocol_var,
                                           values=self.protocol_options, width=7)
        self.protocol_menu.pack(side=tk.LEFT, padx=5, pady=(15, 5))

    def create_buttons(self):
        self.start_button = tk.Button(self.protocol_frame, text="开始嗅探", command=self.start_sniffing, bg="#4CAF50",
                                      fg="white")
        self.start_button.pack(side=tk.LEFT, padx=5, pady=(15, 5))

        self.stop_button = tk.Button(self.protocol_frame, text="停止嗅探", command=self.stop_sniffing, bg="#f44336",
                                     fg="white")
        self.stop_button.pack(side=tk.LEFT, padx=5, pady=(15, 5))
        self.stop_button.config(state=tk.DISABLED)

        self.advanced_button = tk.Button(self.protocol_frame, text="进程追踪", command=self.open_advanced_window, bg="#2196F3", fg="white")
        self.advanced_button.pack(side=tk.LEFT, padx=5, pady=(15, 5))

    def get_nic_list(self):
        return [iface for iface in psutil.net_if_addrs().keys()]

    def create_packet_display_frame(self):
        self.packet_tree_frame = tk.Frame(self.main_frame)
        self.packet_tree_frame.pack(fill=tk.BOTH, expand=True)

        self.packet_tree = ttk.Treeview(self.packet_tree_frame, columns=(
            "Time", "Source", "Dest", "Protocol", "Length", "Source Port", "Dest Port"), show="headings")
        self.packet_tree.heading("Time", text="时间", anchor=tk.W)
        self.packet_tree.heading("Source", text="源", anchor=tk.W)
        self.packet_tree.heading("Dest", text="目的", anchor=tk.W)
        self.packet_tree.heading("Protocol", text="协议", anchor=tk.W)
        self.packet_tree.heading("Length", text="长度", anchor=tk.W)
        self.packet_tree.heading("Source Port", text="源端口", anchor=tk.W)
        self.packet_tree.heading("Dest Port", text="目的端口", anchor=tk.W)

        for col in ("Time", "Source", "Dest", "Protocol", "Length", "Source Port", "Dest Port"):
            self.packet_tree.column(col, width=120, anchor=tk.W)

        self.v_scrollbar = tk.Scrollbar(self.packet_tree_frame, orient=tk.VERTICAL, command=self.packet_tree.yview)
        self.v_scrollbar.pack(side=tk.RIGHT, fill=tk.Y)

        self.packet_tree.configure(yscrollcommand=self.v_scrollbar.set)
        self.packet_tree.pack(padx=10, pady=7, fill=tk.BOTH, expand=True)

        self.packet_tree.bind("<ButtonRelease-1>", self.on_packet_select)

    def create_packet_details_frame(self):
        self.packet_details = scrolledtext.ScrolledText(self.main_frame, width=60, height=10, bg="#FFFFFF",
                                                        font=("Arial", 10))
        self.packet_details.pack(padx=10, pady=7, fill=tk.BOTH, expand=True)

    def create_16_data_frame(self):
        self.data = scrolledtext.ScrolledText(self.main_frame, width=60, height=9, bg="#FFFFFF",
                                                     font=("Courier New", 10))
        self.data.pack(padx=10, pady=7, fill=tk.BOTH, expand=True)

    def start_sniffing(self):
        self.packet_tree.delete(*self.packet_tree.get_children())
        self.captured_packets.clear()
        self.sniffing = True
        self.stop_button.config(state=tk.NORMAL)
        self.sniffer_thread = threading.Thread(target=self.sniff_packets)
        self.sniffer_thread.daemon = True
        self.sniffer_thread.start()

    def stop_sniffing(self):
        self.sniffing = False
        self.stop_button.config(state=tk.DISABLED)
        self.show_statistics()

    def stop_sniffing2(self):
        self.sniffing = False
        self.stop_button.config(state=tk.DISABLED)

    def show_statistics(self):
        stats_window = tk.Toplevel(self.root)
        stats_window.title("数据包统计")
        stats_window.geometry("400x300")

        ip_count = {}
        for packet in self.captured_packets:
            if IP in packet:
                ip_src = packet[IP].src
                ip_dst = packet[IP].dst
                ip_count[ip_src] = ip_count.get(ip_src, 0) + 1
                ip_count[ip_dst] = ip_count.get(ip_dst, 0) + 1

        sorted_ips = sorted(ip_count.items(), key=lambda x: x[1], reverse=True)[:5]

        # 创建表格
        stats_frame = tk.Frame(stats_window)
        stats_frame.pack(pady=20)

        header_label = tk.Label(stats_frame, text="较活跃的前5个IP地址", font=("Arial", 13))
        header_label.grid(row=0, column=0, columnspan=2)

        tk.Label(stats_frame, text="IP地址", font=("Arial", 11)).grid(row=1, column=0, padx=10, pady=5)
        tk.Label(stats_frame, text="次数", font=("Arial", 11)).grid(row=1, column=1, padx=10, pady=5)

        # 填充表格
        for index, (ip, count) in enumerate(sorted_ips, start=2):
            tk.Label(stats_frame, text=ip).grid(row=index, column=0, padx=10, pady=5)
            tk.Label(stats_frame, text=count).grid(row=index, column=1, padx=10, pady=5)

    def sniff_packets(self):
        nic = self.nic_var.get()
        filter_str = ""
        if self.protocol_var.get() == "HTTP":
            filter_str = "tcp port 80"
        elif self.protocol_var.get() == "ICMP":
            filter_str = "icmp"
        elif self.protocol_var.get() == "TCP":
            filter_str = "tcp"
        elif self.protocol_var.get() == "UDP":
            filter_str = "udp"

        sniff(iface=nic, prn=self.process_packet, filter=filter_str, store=0)

    def process_packet(self, packet):
        if not self.sniffing:
            return

        protocol = "Unknown"
        source_port = dest_port = "N/A"

        if IP in packet:
            if TCP in packet:
                protocol = "TCP"
                source_port = packet[TCP].sport
                dest_port = packet[TCP].dport
            elif UDP in packet:
                protocol = "UDP"
                source_port = packet[UDP].sport
                dest_port = packet[UDP].dport
            elif ICMP in packet:
                protocol = "ICMP"
                dest_port = "N/A"  # ICMP doesn't use ports

            # 处理HTTP数据包
            if Raw in packet and b"HTTP" in packet[Raw].load:
                protocol = "HTTP"
                source_port = packet[TCP].sport if TCP in packet else source_port
                dest_port = packet[TCP].dport if TCP in packet else dest_port

            # 筛选IP和端口
            host_ip = self.host_var.get().strip()
            port = self.port_var.get().strip()

            ip_matches = (host_ip == packet[IP].src or host_ip == packet[IP].dst) if host_ip else True
            port_matches = (source_port == int(port) or dest_port == int(port)) if port and port.isdigit() else True

            # 更新协议筛选条件，确保HTTP和ICMP被包含
            protocol_matches = (self.protocol_var.get() == "ALL" or
                                protocol in self.protocol_var.get() or
                                (protocol == "HTTP" and "HTTP" in self.protocol_var.get()) or
                                (protocol == "ICMP" and "ICMP" in self.protocol_var.get()))

            if protocol_matches and ip_matches and port_matches:
                self.add_packet_to_tree(packet, protocol, source_port, dest_port)

    def process_packet_for_selected_process(self, packet):
        if not self.sniffing:
            return

        protocol = "Unknown"
        source_port = dest_port = "N/A"

        if IP in packet:
            if TCP in packet:
                protocol = "TCP"
                source_port = packet[TCP].sport
                dest_port = packet[TCP].dport
            elif UDP in packet:
                protocol = "UDP"
                source_port = packet[UDP].sport
                dest_port = packet[UDP].dport
            elif ICMP in packet:
                protocol = "ICMP"

            # 判断包的源和目的进程是否与选定的PID匹配
            for conn in psutil.net_connections(kind='inet'):
                if conn.pid == self.selected_pid:
                    if (conn.laddr.port == source_port or conn.raddr.port == dest_port):
                        # 只有当协议匹配时才添加到树中
                        if (self.protocol_var.get() == "ALL" or
                                protocol in self.protocol_var.get()):
                            self.add_packet_to_tree(packet, protocol, source_port, dest_port)
                        break

    def add_packet_to_tree(self, packet, protocol, source_port, dest_port):
        if IP in packet:
            ip_layer = packet[IP]
            timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
            length = len(packet)
            self.packet_tree.insert("", "end", values=(
                timestamp, ip_layer.src, ip_layer.dst, protocol, length, source_port, dest_port))
            self.captured_packets.append(packet)

    def on_packet_select(self, event):
        selected_item = self.packet_tree.selection()
        if selected_item:
            index = self.packet_tree.index(selected_item[0])
            packet = self.captured_packets[index]
            self.packet_details.delete(1.0, tk.END)
            details = self.get_packet_info(packet)
            self.packet_details.insert(tk.END, details)
            self.show_16_data(packet)

    def get_packet_info(self, packet):
        return packet.show(dump=True)

    def show_16_data(self, packet):
        self.data.delete(1.0, tk.END)
        self.data.insert(tk.END, hexdump(packet, dump=True))

    def open_advanced_window(self):
        self.advanced_window = tk.Toplevel(self.root)
        self.advanced_window.title("进程追踪")
        self.advanced_window.geometry("400x300")
        self.advanced_window.grab_set()

        self.process_listbox = tk.Listbox(self.advanced_window, selectmode=tk.SINGLE)
        self.process_listbox.pack(padx=10, pady=10, fill=tk.BOTH, expand=True)

        self.populate_process_list()

        self.select_button = tk.Button(self.advanced_window, text="选择进程", command=self.select_process)
        self.select_button.pack(pady=(5, 10))

    def populate_process_list(self):
        for proc in psutil.process_iter(attrs=['pid', 'name']):
            self.process_listbox.insert(tk.END, f"{proc.info['pid']} - {proc.info['name']}")

    def select_process(self):
        selection = self.process_listbox.curselection()
        if selection:
            selected_process = self.process_listbox.get(selection[0])
            self.selected_pid = int(selected_process.split(" - ")[0])
            messagebox.showinfo("选择成功", f"已选择进程: {selected_process}")
            self.start_sniffing_for_process()

    def start_sniffing_for_process(self):
        self.stop_sniffing2()  # 停止当前嗅探
        self.packet_tree.delete(*self.packet_tree.get_children())  # 清空当前包列表
        self.captured_packets.clear()  # 清空捕获包列表

        # 重新开始嗅探，仅嗅探所选进程的TCP数据包
        self.sniffing = True
        self.stop_button.config(state=tk.NORMAL)
        self.sniffer_thread = threading.Thread(target=self.sniff_packets_for_selected_process)
        self.sniffer_thread.daemon = True
        self.sniffer_thread.start()

    def sniff_packets_for_selected_process(self):
        nic = self.nic_var.get()
        sniff(iface=nic, prn=self.process_packet_for_selected_process, store=0)

    def process_packet_for_selected_process(self, packet):
        if not self.sniffing:
            return

        protocol = "Unknown"
        source_port = dest_port = "N/A"

        if IP in packet:
            if TCP in packet:
                protocol = "TCP"
                source_port = packet[TCP].sport
                dest_port = packet[TCP].dport
            elif UDP in packet:
                protocol = "UDP"
                source_port = packet[UDP].sport
                dest_port = packet[UDP].dport

            # 判断包的源和目的进程是否与选定的PID匹配
            for conn in psutil.net_connections(kind='inet'):
                if conn.pid == self.selected_pid:
                    if conn.laddr.port == source_port or conn.raddr.port == dest_port:
                        self.add_packet_to_tree(packet, protocol, source_port, dest_port)
                        break

    def run(self):
        self.root.mainloop()
